"""
Dify API 简化工具类
作者：youxin

这个工具类提供了简洁的 Dify API 集成功能，专注于：
- 非流式和流式请求处理
- 完整的错误处理和异常管理
- 返回原始数据，由用户自行处理数据组合和保存

设计理念：
- 保持简单，专注核心功能
- 错误处理完善，便于调试
- 返回原始数据，用户自由处理
"""

import requests
import json
from typing import Dict, Any, Iterator, Optional


# ==================== 异常处理类 ====================

class DifyAPIError(Exception):
    """Dify API 基础异常类"""
    
    def __init__(self, message: str, status_code: int = None, error_code: str = None):
        """
        初始化异常
        
        Args:
            message: 错误消息
            status_code: HTTP状态码
            error_code: Dify错误码
        """
        self.message = message
        self.status_code = status_code
        self.error_code = error_code
        super().__init__(self.message)
    
    def __str__(self):
        """返回格式化的错误信息"""
        if self.status_code and self.error_code:
            return f"[{self.status_code}] {self.error_code}: {self.message}"
        elif self.status_code:
            return f"[{self.status_code}] {self.message}"
        else:
            return self.message


class InvalidParamError(DifyAPIError):
    """400 - 传入参数异常"""
    pass


class AppUnavailableError(DifyAPIError):
    """400 - App 配置不可用"""
    pass


class ProviderNotInitializeError(DifyAPIError):
    """400 - 无可用模型凭据配置"""
    pass


class ProviderQuotaExceededError(DifyAPIError):
    """400 - 模型调用额度不足"""
    pass


class ModelCurrentlyNotSupportError(DifyAPIError):
    """400 - 当前模型不可用"""
    pass


class WorkflowRequestError(DifyAPIError):
    """400 - workflow 执行失败"""
    pass


class InternalServerError(DifyAPIError):
    """500 - 服务内部异常"""
    pass


# ==================== 主工具类 ====================

class DifyWorkflowToolkit:
    """
    Dify 工作流简化工具类
    
    专注于提供简洁的API调用接口，返回原始数据供用户自行处理
    """
    
    def __init__(self, api_key: str, user_id: str, base_url: str = "http://dify.dajiaok.com/v1"):
        """
        初始化 Dify 工作流工具类
        
        Args:
            api_key: Dify API 密钥
            user_id: 用户ID，用于标识请求来源
            base_url: API 基础URL，默认为官方地址
        """
        self.base_url = base_url.rstrip('/')  # 移除末尾斜杠
        self.api_key = api_key
        self.user_id = user_id
        
        # 设置请求头
        self.headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json;charset=UTF-8",
            "Accept-Charset": "UTF-8"
        }
    
    def run_workflow_blocking(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """
        运行工作流（非流式模式）
        
        发送请求并等待完整响应，适用于需要一次性获取完整结果的场景
        
        Args:
            inputs: 工作流输入参数字典
            
        Returns:
            Dict[str, Any]: 原始响应数据，包含以下字段：
                - workflow_run_id: 工作流运行ID
                - task_id: 任务ID
                - data: 工作流执行结果数据
                    - id: 执行实例ID
                    - workflow_id: 工作流ID
                    - status: 执行状态 (completed/failed)
                    - outputs: 输出结果
                    - error: 错误信息（如果有）
                    - elapsed_time: 执行耗时（秒）
                    - total_tokens: 使用的token总数
                    - total_steps: 执行步骤总数
                    - created_at: 创建时间戳
                    - finished_at: 完成时间戳
                    
        Raises:
            DifyAPIError: API调用失败时抛出相应的异常
            
        Example:
            >>> toolkit = DifyWorkflowToolkit(api_key="your_key", user_id="user123")
            >>> result = toolkit.run_workflow_blocking({"query": "你好"})
            >>> print(result['data']['outputs'])
        """
        url = f"{self.base_url}/workflows/run"
        
        # 构建请求载荷
        payload = {
            "inputs": inputs,
            "user": self.user_id,
            "response_mode": "blocking"  # 非流式模式
        }
        
        try:
            # 发送POST请求
            response = requests.post(url, headers=self.headers, json=payload)
            
            # 检查响应状态
            if not response.ok:
                raise self._parse_error_response(response)
            
            # 返回原始JSON数据
            return response.json()
            
        except requests.RequestException as e:
            # 网络请求异常
            raise DifyAPIError(f"网络请求失败: {str(e)}")
    
    def run_workflow_streaming(self, inputs: Dict[str, Any]) -> Iterator[Dict[str, Any]]:
        """
        运行工作流（流式模式）
        
        发送流式请求并返回数据迭代器，适用于需要实时获取执行过程的场景
        
        Args:
            inputs: 工作流输入参数字典
            
        Yields:
            Dict[str, Any]: 流式响应数据块，每个数据块包含：
                - event: 事件类型
                    - workflow_started: 工作流开始
                    - node_started: 节点开始执行
                    - text_chunk: 文本块输出
                    - node_finished: 节点执行完成
                    - workflow_finished: 工作流完成
                    - tts_message: TTS音频消息
                    - tts_message_end: TTS音频结束
                    - ping: 保持连接
                - workflow_run_id: 工作流运行ID
                - task_id: 任务ID
                - data: 事件相关数据
                - message_id: 消息ID（可选）
                
        Raises:
            DifyAPIError: API调用失败时抛出相应的异常
            
        Example:
            >>> toolkit = DifyWorkflowToolkit(api_key="your_key", user_id="user123")
            >>> for chunk in toolkit.run_workflow_streaming({"query": "你好"}):
            ...     if chunk['event'] == 'text_chunk':
            ...         print(chunk['data']['text'], end='')
            ...     elif chunk['event'] == 'workflow_finished':
            ...         print(f"\\n工作流完成，状态: {chunk['data']['status']}")
        """
        url = f"{self.base_url}/workflows/run"
        
        # 构建请求载荷
        payload = {
            "inputs": inputs,
            "user": self.user_id,
            "response_mode": "streaming"  # 流式模式
        }
        
        try:
            # 发送流式POST请求
            response = requests.post(url, headers=self.headers, json=payload, stream=True)
            
            # 检查响应状态
            if not response.ok:
                raise self._parse_error_response(response)
            
            # 处理流式响应 使用response.iter_lines()方法逐行迭代响应内容。这个方法会按行分割响应体，并且会等待数据流中的新行。
            for line in response.iter_lines():
                if line:
                    line_str = line.decode('utf-8')
                    
                    # 跳过空行和非数据行
                    if not line_str.strip() or not line_str.startswith('data: '):
                        continue
                    
                    # 提取数据部分（移除 'data: ' 前缀）
                    data_str = line_str[6:]
                    
                    # 跳过结束标志
                    if data_str.strip() == '[DONE]':
                        break
                    
                    try:
                        # 解析JSON数据并返回
                        chunk_data = json.loads(data_str)
                        if isinstance(chunk_data, dict) and 'event' in chunk_data:
                            yield chunk_data
                    except json.JSONDecodeError as e:
                        # JSON解析失败，跳过这一行
                        continue
                        
        except requests.RequestException as e:
            # 网络请求异常
            raise DifyAPIError(f"网络请求失败: {str(e)}")
    
    def _parse_error_response(self, response: requests.Response) -> DifyAPIError:
        """
        解析错误响应并返回对应的异常对象
        
        根据HTTP状态码和Dify错误码返回具体的异常类型，便于用户进行针对性处理
        
        Args:
            response: HTTP响应对象
            
        Returns:
            DifyAPIError: 对应的异常对象
        """
        status_code = response.status_code
        
        try:
            # 尝试解析错误响应的JSON数据
            error_data = response.json()
            error_code = error_data.get('code', 'unknown')
            message = error_data.get('message', f'HTTP {status_code} Error')
        except (json.JSONDecodeError, ValueError):
            # JSON解析失败，使用默认错误信息
            error_code = 'unknown'
            message = f'HTTP {status_code} Error: {response.text[:200]}'
        
        # 根据状态码和错误码返回对应的异常类型
        if status_code == 400:
            # 400错误的具体分类
            error_mapping = {
                'invalid_param': InvalidParamError,
                'app_unavailable': AppUnavailableError,
                'provider_not_initialize': ProviderNotInitializeError,
                'provider_quota_exceeded': ProviderQuotaExceededError,
                'model_currently_not_support': ModelCurrentlyNotSupportError,
                'workflow_request_error': WorkflowRequestError,
            }
            
            exception_class = error_mapping.get(error_code, DifyAPIError)
            return exception_class(message, status_code, error_code)
            
        elif status_code == 500:
            # 500服务器内部错误
            return InternalServerError(message, status_code, error_code)
        else:
            # 其他HTTP错误
            return DifyAPIError(message, status_code, error_code)
    
    def get_api_info(self) -> Dict[str, str]:
        """
        获取当前配置的API信息
        
        Returns:
            Dict[str, str]: 包含API配置信息的字典
        """
        return {
            "base_url": self.base_url,
            "user_id": self.user_id,
            "api_key_prefix": f"{self.api_key[:8]}..." if len(self.api_key) > 8 else "***"
        }
    
    def get_workflow_result(self, workflow_run_id: str) -> Dict[str, Any]:
        """
        获取工作流执行结果
        
        根据工作流执行ID获取当前执行状态和结果，适用于查询已启动工作流的执行情况
        
        Args:
            workflow_run_id: 工作流执行ID，通过 run_workflow_blocking 或 run_workflow_streaming 获得
            
        Returns:
            Dict[str, Any]: 工作流执行结果，包含以下字段：
                - id: 工作流执行ID
                - workflow_id: 关联的工作流ID
                - status: 执行状态 (running/succeeded/failed/stopped)
                - inputs: 任务输入内容 (JSON格式)
                - outputs: 任务输出内容 (JSON格式，执行完成后才有)
                - error: 错误原因 (执行失败时才有)
                - total_steps: 任务执行总步数
                - total_tokens: 任务执行总tokens
                - created_at: 任务开始时间戳
                - finished_at: 任务结束时间戳 (执行完成后才有)
                - elapsed_time: 耗时(秒)
                
        Raises:
            DifyAPIError: API调用失败时抛出相应的异常
            
        Example:
            >>> toolkit = DifyWorkflowToolkit(api_key="your_key", user_id="user123")
            >>> # 先启动一个工作流
            >>> result = toolkit.run_workflow_blocking({"query": "你好"})
            >>> workflow_run_id = result['workflow_run_id']
            >>> 
            >>> # 获取执行结果
            >>> status_result = toolkit.get_workflow_result(workflow_run_id)
            >>> print(f"执行状态: {status_result['status']}")
            >>> if status_result['status'] == 'succeeded':
            ...     print(f"输出结果: {status_result['outputs']}")
            >>> elif status_result['status'] == 'failed':
            ...     print(f"错误信息: {status_result['error']}")
        """
        if not workflow_run_id:
            raise InvalidParamError("workflow_run_id 不能为空")
        
        url = f"{self.base_url}/workflows/run/{workflow_run_id}"
        
        try:
            # 发送GET请求获取工作流执行结果
            response = requests.get(url, headers=self.headers)
            
            # 检查响应状态
            if not response.ok:
                raise self._parse_error_response(response)
            
            # 返回原始JSON数据
            result_data = response.json()
            
            # 处理inputs和outputs字段，如果是字符串则尝试解析为JSON
            if 'inputs' in result_data and isinstance(result_data['inputs'], str):
                try:
                    result_data['inputs'] = json.loads(result_data['inputs'])
                except json.JSONDecodeError:
                    # 如果解析失败，保持原始字符串
                    pass
            
            if 'outputs' in result_data and isinstance(result_data['outputs'], str):
                try:
                    result_data['outputs'] = json.loads(result_data['outputs'])
                except json.JSONDecodeError:
                    # 如果解析失败，保持原始字符串
                    pass
            
            return result_data
            
        except requests.RequestException as e:
            # 网络请求异常
            raise DifyAPIError(f"网络请求失败: {str(e)}")


# ==================== 使用示例 ====================

if __name__ == "__main__":
    """
    简单的使用示例
    """
    
    # 初始化工具类
    toolkit = DifyWorkflowToolkit(
        # api_key 代表工作流的唯一标识符，用于认证和授权
        api_key="***",
        user_id="***"
    )
    
    api_info = toolkit.get_api_info()
    print("🔧 当前API配置:")
    for key, value in api_info.items():
        print(f"   {key}: {value}")

    # 输入参数 根据工作流实际定义参数调整 这里的input参数根据工作流定义的输入参数调整
    inputs = {"input": "你好，我是一名Java小白，请你帮我规划一套Java的学习计划"}
    
    print("=== 非流式模式示例 ===")
    try:
        # 非流式调用
        result = toolkit.run_workflow_blocking(inputs)
        print(f"工作流ID: {result['workflow_run_id']}")
        print(f"执行状态: {result['data']['status']}")
        print(f"输出结果: {result['data']['outputs']['text']}")
        with open("./run_workflow_blocking.html", "w", encoding="utf-8") as f:
            f.write(result['data']['outputs']['text'])
    except DifyAPIError as e:
        print(f"API调用失败: {e}")
    
    print("\n=== 流式模式示例 ===")
    try:
        # 流式调用
        collected_text = []
        for chunk in toolkit.run_workflow_streaming(inputs):
            event = chunk['event']
            data = chunk['data']
            # 简单调用，详情请看 readme.md 中的 流式调用 部分
            if event == 'workflow_started':
                print(f"🚀 工作流开始: {data.get('workflow_id')}")
            elif event == 'text_chunk':
                text = data.get('text', '')
                collected_text.append(text)
                print(text, end='', flush=True)
            elif event == 'workflow_finished':
                print(f"\n✅ 工作流完成，状态: {data.get('status')}")
                print(f"总耗时: {data.get('elapsed_time', 0):.2f}秒")
                print(f"使用Token: {data.get('total_tokens', 0)}")
        
        # 用户可以自行处理收集到的文本
        full_text = ''.join(collected_text)
        print(f"\n完整文本长度: {len(full_text)} 字符")
        with open("./run_workflow_streaming.html", "w", encoding="utf-8") as f:
            f.write(full_text)
    except DifyAPIError as e:
        print(f"流式调用失败: {e}")
    
    print("\n=== 获取工作流执行结果示例 ===")
    try:
        # 先启动一个非流式工作流获取workflow_run_id
        workflow_run_id = "***"
        
        # 获取工作流执行结果
        status_result = toolkit.get_workflow_result(workflow_run_id)
        print(f"🔍 工作流执行状态: {status_result['status']}")
        print(f"📊 执行步骤数: {status_result['total_steps']}")
        print(f"🔢 使用Token数: {status_result['total_tokens']}")
        print(f"⏱️ 执行耗时: {status_result['elapsed_time']:.2f}秒")
        
        if status_result['status'] == 'succeeded':
            print("✅ 工作流执行成功")
            if status_result.get('outputs'):
                print(f"📤 输出结果类型: {type(status_result['outputs'])}")
                print(f"📤 输出结果: {status_result['outputs']['text']}")
        elif status_result['status'] == 'failed':
            print(f"❌ 工作流执行失败: {status_result.get('error', '未知错误')}")
        elif status_result['status'] == 'running':
            print("🔄 工作流正在执行中...")
        
    except DifyAPIError as e:
        print(f"获取工作流结果失败: {e}")